# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import math

import coin.exchange.base.kr_rest.private_client_base as privcb
import coin.exchange.bitmex.kr_rest.native_private_client as nprivc
import coin.exchange.bitmex.kr_rest.futures_product as exchange_product
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate

from coin.exchange.base.order_enum import convert_two_order_side
import coin.proto.coin_order_enums_pb2 as coin_order
from coin.proto.coin_query_pb2 import (CurrencyBalance,
                                       AccountBalance,
                                       MarginUsage,
                                       AccountPosition,
                                       ProductPosition,
                                       ProductFillElement,
                                       ProductFillBundle,
                                       ProductOrderElement,
                                       ProductOrderBundle)
from coin.proto.coin_order_gateway_pb2 import (
    OrderSubmitResponseProto,
    OrderCancelResponseProto,
)


class BitmexPrivateParser(object):
  @staticmethod
  def parse_balance(json_msg):
    assert json_msg['currency'] == 'XBt'
    # Convert from satoshi to BTC.
    position_margin = float(json_msg['maintMargin']) / 1e8
    order_margin = float(json_msg['initMargin']) / 1e8
    assert position_margin >= 0.0
    assert order_margin >= 0.0
    total = json_msg['marginBalance'] / 1e8
    hold = position_margin + order_margin
    available = float(json_msg['availableMargin']) / 1e8
    assert math.isclose(total, hold + available)
    assert available >= 0.0
    currency_balance = CurrencyBalance(
        currency='BTC',
        currency_native='XBT',
        total=total,
        hold=hold,
        available=available,
    )
    return AccountBalance(exchange="Bitmex", market_type="Futures", each_balance=[currency_balance])

  @staticmethod
  def parse_margin(json_msg):
    assert json_msg['currency'] == 'XBt'
    # Convert from satoshi to BTC.
    position_margin = float(json_msg['maintMargin']) / 1e8
    order_margin = float(json_msg['initMargin']) / 1e8
    assert position_margin >= 0.0
    assert order_margin >= 0.0
    total_margin = json_msg['marginBalance'] / 1e8
    hold_margin = position_margin + order_margin
    available_margin = float(json_msg['availableMargin']) / 1e8
    assert math.isclose(total_margin, hold_margin + available_margin)
    assert available_margin >= 0.0

    total_margin = float(json_msg['marginBalance']) / 1e8
    return MarginUsage(
        currency='BTC',
        total_margin=total_margin,
        hold_margin=hold_margin,
        available_margin=available_margin,
    )

  @staticmethod
  def parse_submit_response(json_msg, product):
    # If there is error field in the message, then some error happened.
    if json_msg.get('error'):
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          symbol=product.symbol,
          success=False,
      )
    else:  # Else order submission is considered successful.
      assert json_msg['symbol'] == product.native_symbol
      order_id = json_msg['orderID']
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          symbol=product.symbol,
          success=True,
          order_id=order_id,
      )

  @staticmethod
  def parse_cancel_response(json_msg, order_id):
    # If there is error field in the message, then some error happened.
    if isinstance(json_msg, dict) and json_msg.get('error'):
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          symbol=product.symbol,
          success=False,
      )
    elif isinstance(json_msg, list):
      assert json_msg[0]['orderID'] == order_id
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          success=True,
          order_id=order_id,
      )
    else:
      raise ValueError('Invalid input %s' % json_msg)

  @staticmethod
  def parse_cancel_all_response(json_msg):
    # If there is error field in the message, then some error happened.
    if isinstance(json_msg, dict) and json_msg.get('error'):
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          symbol=product.symbol,
          success=False,
      )
    else:  # Else order cancellation is considered successful.
      return OrderSubmitResponseProto(
          exchange='Bitmex',
          success=True,
      )

  @staticmethod
  def parse_position(json_msg: list, verbose=False):
    assert isinstance(json_msg, list)
    position_list = []
    for msg_elem in json_msg:
      product = BitmexFuturesProduct.FromStr(msg_elem['symbol'])
      net_position = msg_elem['currentQty']
      realized_pnl = msg_elem['realisedPnl'] / 1e8
      unrealized_pnl = msg_elem['unrealisedPnl'] / 1e8
      total_pnl = realized_pnl + unrealized_pnl
      liquidation_price = msg_elem['liquidationPrice']

      position = ProductPosition(symbol=product.symbol,
                                 native_symbol=product.native_symbol,
                                 net_position=net_position,
                                 realized_pnl=realized_pnl,
                                 unrealized_pnl=unrealized_pnl,
                                 total_pnl=total_pnl,
                                 liquidation_price=liquidation_price)
      position_list.append(position)
    return position_list

  @staticmethod
  def parse_order_history(json_msg: list):
    future_fill_list = []
    open_order_list = []
    for msg_elem in json_msg:
      if msg_elem['side'] == 'Sell':
        side = coin_order.SELL_ORDER
      elif msg_elem['side'] == 'Buy':
        side = coin_order.BUY_ORDER
      else:
        raise ValueError(msg_elem['side'])

      order_qty = float(msg_elem['orderQty'])
      leaves_qty = float(msg_elem['leavesQty'])
      assert order_qty >= 0
      assert 0 <= leaves_qty <= order_qty, (order_qty, leaves_qty)
      price = float(msg_elem['price'])
      assert price >= 0.0
      fill_qty = order_qty - leaves_qty
      native_symbol = msg_elem['symbol']
      product = BitmexFuturesProduct.FromStr(native_symbol)

      if math.isclose(fill_qty, 0.0):  # New order
        open_order = ProductOrderElement(symbol=product.symbol,
                                         side=side,
                                         price=price,
                                         qty=order_qty,
                                         order_id=msg_elem['orderID'])
        open_order_list.append(open_order)
      elif math.isclose(leaves_qty, 0.0):  # Filled order
        future_fill = ProductFillElement(symbol=product.symbol,
                                         side=side,
                                         price=price,
                                         qty=fill_qty,
                                         order_id=msg_elem['orderID'])
        future_fill_list.append(future_fill)
      else:  # Partially filled.
        open_order = ProductOrderElement(symbol=product.symbol,
                                         side=side,
                                         price=price,
                                         qty=order_qty,
                                         order_id=msg_elem['orderID'])
        open_order_list.append(open_order)

        future_fill = ProductFillElement(symbol=product.symbol,
                                         side=side,
                                         price=price,
                                         qty=fill_qty,
                                         order_id=msg_elem['orderID'])
        future_fill_list.append(future_fill)

    fill_proto = ProductFillBundle(
        exchange='Bitmex',
        each_fill=future_fill_list,
    )
    open_order_proto = ProductOrderBundle(
        exchange='Bitmex',
        each_order=open_order_list,
    )
    return fill_proto, open_order_proto


class BitmexPrivateClient(privcb.PrivateClientBase):
  ProductType = exchange_product.BitmexFuturesProduct

  def __init__(self, key_file, timeout=5, test=False):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.BitmexNativePrivateClient(key_file=key_file, timeout=timeout, test=test)

  def query_account_balance_impl(self):
    update = self.nprivc.query_user_margin()
    update.msg = BitmexPrivateParser.parse_balance(update.msg)
    return update

  def query_account_position_impl(self):
    update_margin = self.nprivc.query_user_margin()
    update_position = self.nprivc.query_position()
    update_margin.msg = BitmexPrivateParser.parse_margin(update_margin.msg)
    position_list = BitmexPrivateParser.parse_position(update_position.msg)
    account_position = AccountPosition(
        exchange="Bitmex",
        market_type="Futures",
        each_position=position_list,
        each_margin_usage=[update_margin.msg],
    )
    update = RestUpdate(
        req_timestamp=update_margin.req_timestamp,
        res_timestamp=update_position.res_timestamp,
        msg=account_position,
        native_msg={
            "margin": update_margin.native_msg,
            "position": update_position.native_msg,
        },
    )
    return update

  def query_fills_impl(self, product):
    # TODO(xguo) Is there a partial-fill status?
    update = self.nprivc.query_orders(product.native_symbol, status=None)
    # Ignore open orders
    fill_proto, _ = \
        BitmexPrivateParser.parse_order_history(update.msg)
    fill_proto.symbol = product.symbol
    update.msg = fill_proto
    return update

  def query_account_fills_impl(self):
    update = self.nprivc.query_orders(symbol=None, status=None)
    # Ignore open orders
    fill_proto, _ = \
        BitmexPrivateParser.parse_order_history(update.msg)
    update.msg = fill_proto
    return update

  def query_list_orders_impl(self, product):
    update = self.nprivc.query_orders(product.native_symbol, status='New')
    _, open_orders_proto = \
        BitmexPrivateParser.parse_order_history(update.msg)
    open_orders_proto.symbol = product.symbol
    update.msg = open_orders_proto
    return update

  def query_account_list_orders_impl(self):
    update = self.nprivc.query_orders(symbol=None)
    _, open_orders_proto = \
        BitmexPrivateParser.parse_order_history(update.msg)
    return open_orders_proto

  def query_submit_impl(self, product, side, price, qty):
    side = convert_two_order_side(side)
    side = side.capitalize()
    update = self.nprivc.query_order_new(product.native_symbol, side, price, qty)
    update.msg = BitmexPrivateParser.parse_submit_response(update.msg, product)
    return update

  def query_cancel_impl(self, product, order_id):
    update = self.nprivc.query_order_cancel(order_id)
    update.msg = BitmexPrivateParser.parse_cancel_response(update.msg, order_id)
    return update

  def query_cancel_product_impl(self, product):
    update = self.nprivc.query_order_cancel_all(product.native_symbol)
    update.msg = BitmexPrivateParser.parse_cancel_all_response(update.msg)
    return update
